{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [],
   "source": [
    "import boto3\n",
    "import sagemaker\n",
    "import pandas as pd\n",
    "\n",
    "sess   = sagemaker.Session()\n",
    "bucket = sess.default_bucket()\n",
    "role = sagemaker.get_execution_role()\n",
    "region = boto3.Session().region_name\n",
    "\n",
    "sm = boto3.Session().client(service_name='sagemaker', region_name=region)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Specify the S3 Location of the Features"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "no stored variable scikit_processing_job_s3_output_prefix\n"
     ]
    }
   ],
   "source": [
    "%store -r scikit_processing_job_s3_output_prefix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Previous Scikit Processing Job Name: sagemaker-scikit-learn-2020-03-30-03-34-18-188\n"
     ]
    }
   ],
   "source": [
    "# HACK\n",
    "scikit_processing_job_s3_output_prefix = 'sagemaker-scikit-learn-2020-03-30-03-34-18-188'\n",
    "\n",
    "print('Previous Scikit Processing Job Name: {}'.format(scikit_processing_job_s3_output_prefix))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "metadata": {},
   "outputs": [],
   "source": [
    "prefix_train = '{}/output/bert-train'.format(scikit_processing_job_s3_output_prefix)\n",
    "prefix_validation = '{}/output/bert-validation'.format(scikit_processing_job_s3_output_prefix)\n",
    "prefix_test = '{}/output/bert-test'.format(scikit_processing_job_s3_output_prefix)\n",
    "\n",
    "path_train = './{}'.format(prefix_train)\n",
    "path_validation = './{}'.format(prefix_validation)\n",
    "path_test = './{}'.format(prefix_test)\n",
    "\n",
    "train_s3_uri = 's3://{}/{}'.format(bucket, prefix_train)\n",
    "validation_s3_uri = 's3://{}/{}'.format(bucket, prefix_validation)\n",
    "test_s3_uri = 's3://{}/{}'.format(bucket, prefix_test)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'DataSource': {'S3DataSource': {'S3DataType': 'S3Prefix', 'S3Uri': 's3://sagemaker-us-east-1-835319576252/sagemaker-scikit-learn-2020-03-30-03-34-18-188/output/bert-train', 'S3DataDistributionType': 'ShardedByS3Key'}}}\n",
      "{'DataSource': {'S3DataSource': {'S3DataType': 'S3Prefix', 'S3Uri': 's3://sagemaker-us-east-1-835319576252/sagemaker-scikit-learn-2020-03-30-03-34-18-188/output/bert-validation', 'S3DataDistributionType': 'ShardedByS3Key'}}}\n",
      "{'DataSource': {'S3DataSource': {'S3DataType': 'S3Prefix', 'S3Uri': 's3://sagemaker-us-east-1-835319576252/sagemaker-scikit-learn-2020-03-30-03-34-18-188/output/bert-test', 'S3DataDistributionType': 'ShardedByS3Key'}}, 'ContentType': 'text/csv'}\n"
     ]
    }
   ],
   "source": [
    "s3_input_train_data = sagemaker.s3_input(s3_data=train_s3_uri, distribution='ShardedByS3Key') #, content_type='text/csv')\n",
    "s3_input_validation_data = sagemaker.s3_input(s3_data=validation_s3_uri, distribution='ShardedByS3Key') #, content_type='text/csv')\n",
    "s3_input_test_data = sagemaker.s3_input(s3_data=test_s3_uri, content_type='text/csv', distribution='ShardedByS3Key') #, content_type='text/csv')\n",
    "\n",
    "print(s3_input_train_data.config)\n",
    "print(s3_input_validation_data.config)\n",
    "print(s3_input_test_data.config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "from sklearn.model_selection import train_test_split\r\n",
      "from sklearn.utils import resample\r\n",
      "\r\n",
      "import pandas as pd\r\n",
      "from datetime import datetime\r\n",
      "import subprocess\r\n",
      "import sys\r\n",
      "\r\n",
      "# We should remove this once the bug is fixed.\r\n",
      "import subprocess\r\n",
      "import sys\r\n",
      "#subprocess.check_call([sys.executable, '-m', 'pip', 'install', 'tensorflow==1.15.2'])\r\n",
      "subprocess.check_call([sys.executable, '-m', 'pip', 'install', 'tensorflow-hub==0.7.0'])\r\n",
      "subprocess.check_call([sys.executable, '-m', 'pip', 'install', 'bert-tensorflow==1.0.1'])\r\n",
      "subprocess.check_call([sys.executable, '-m', 'pip', 'install', 'sagemaker-tensorflow==1.15.0.1.1.0'])\r\n",
      "\r\n",
      "import tensorflow as tf\r\n",
      "print(tf.__version__)\r\n",
      "import tensorflow_hub as hub\r\n",
      "import amazon_run_classifier\r\n",
      "#from bert import run_classifier\r\n",
      "from bert import optimization\r\n",
      "from bert import tokenization\r\n",
      "from tensorflow import keras\r\n",
      "\r\n",
      "import os\r\n",
      "import re\r\n",
      "\r\n",
      "import argparse\r\n",
      "import json\r\n",
      "import os\r\n",
      "import pandas as pd\r\n",
      "import csv\r\n",
      "import glob\r\n",
      "from pathlib import Path\r\n",
      "\r\n",
      "# Based on this...\r\n",
      "# https://github.com/google-research/bert/blob/eedf5716ce1268e56f0a50264a88cafad334ac61/run_classifier.py#L479\r\n",
      "\r\n",
      "MAX_SEQ_LENGTH = 128\r\n",
      "LABEL_VALUES = ['1', '2', '3', '4', '5']\r\n",
      "\r\n",
      "# TODO:  Pass this into the processor\r\n",
      "BERT_MODEL_HUB = \"https://tfhub.dev/google/bert_uncased_L-12_H-768_A-12/1\"\r\n",
      "\r\n",
      "def create_model(is_predicting, \r\n",
      "                 input_ids, \r\n",
      "                 input_mask, \r\n",
      "                 segment_ids, \r\n",
      "                 labels,\r\n",
      "                 num_labels):\r\n",
      "    \"\"\"Creates a classification model.\"\"\"\r\n",
      "\r\n",
      "    bert_module = hub.Module(\r\n",
      "      BERT_MODEL_HUB,\r\n",
      "      trainable=True)\r\n",
      "    bert_inputs = dict(\r\n",
      "      input_ids=input_ids,\r\n",
      "      input_mask=input_mask,\r\n",
      "      segment_ids=segment_ids)\r\n",
      "    bert_outputs = bert_module(\r\n",
      "      inputs=bert_inputs,\r\n",
      "      signature=\"tokens\",\r\n",
      "      as_dict=True)\r\n",
      "\r\n",
      "    # Use \"pooled_output\" for classification tasks on an entire sentence.\r\n",
      "    # Use \"sequence_outputs\" for token-level output.\r\n",
      "    output_layer = bert_outputs[\"pooled_output\"]\r\n",
      "\r\n",
      "    hidden_size = output_layer.shape[-1].value\r\n",
      "\r\n",
      "    # Create our own layer to tune for politeness data.\r\n",
      "    output_weights = tf.get_variable(\r\n",
      "      \"output_weights\", [num_labels, hidden_size],\r\n",
      "      initializer=tf.truncated_normal_initializer(stddev=0.02))\r\n",
      "\r\n",
      "    output_bias = tf.get_variable(\r\n",
      "      \"output_bias\", [num_labels], initializer=tf.zeros_initializer())\r\n",
      "\r\n",
      "    with tf.variable_scope(\"loss\"):\r\n",
      "        # Dropout helps prevent overfitting\r\n",
      "        output_layer = tf.nn.dropout(output_layer, keep_prob=0.9)\r\n",
      "\r\n",
      "        logits = tf.matmul(output_layer, output_weights, transpose_b=True)\r\n",
      "        logits = tf.nn.bias_add(logits, output_bias)\r\n",
      "        log_probs = tf.nn.log_softmax(logits, axis=-1)\r\n",
      "\r\n",
      "        # Convert labels into one-hot encoding\r\n",
      "        print('**** ONE HOT ENCODING THESE LABELS {}'.format(labels))\r\n",
      "        one_hot_labels = tf.one_hot(labels, depth=num_labels, dtype=tf.float32)\r\n",
      "        print('**** ONE HOT ENCODINGS {}'.format(one_hot_labels))\r\n",
      "        \r\n",
      "        print('**** LOG_PROBS {} ****'.format(log_probs))\r\n",
      "        predicted_labels = tf.squeeze(tf.argmax(log_probs, axis=-1, output_type=tf.int32))\r\n",
      "        print('**** PREDICTED LABELS {}'.format(predicted_labels))\r\n",
      "\r\n",
      "        # If we're predicting, we want predicted labels and the probabiltiies.\r\n",
      "        if is_predicting:\r\n",
      "            return (predicted_labels, log_probs)\r\n",
      "\r\n",
      "        # If we're train/eval, compute loss between predicted and actual label\r\n",
      "        per_example_loss = -tf.reduce_sum(one_hot_labels * log_probs, axis=-1)\r\n",
      "        loss = tf.reduce_mean(per_example_loss)\r\n",
      "        return (loss, predicted_labels, log_probs)\r\n",
      "\r\n",
      "\r\n",
      "# model_fn_builder actually creates our model function\r\n",
      "# using the passed parameters for num_labels, learning_rate, etc.\r\n",
      "# This function wraps our model function in a `model_fn_builder` function that adapts our model to work for training, evaluation, and prediction.\r\n",
      "def model_fn_builder(num_labels, learning_rate, num_train_steps,\r\n",
      "                     num_warmup_steps):\r\n",
      "  def model_fn(features, labels, mode, params):\r\n",
      "    input_ids = features[\"input_ids\"]\r\n",
      "    input_mask = features[\"input_mask\"]\r\n",
      "    segment_ids = features[\"segment_ids\"]\r\n",
      "    label_ids = features[\"label_ids\"]\r\n",
      "\r\n",
      "    is_predicting = (mode == tf.estimator.ModeKeys.PREDICT)\r\n",
      "    \r\n",
      "    # TRAIN and EVAL\r\n",
      "    if not is_predicting:\r\n",
      "\r\n",
      "      (loss, predicted_labels, log_probs) = create_model(\r\n",
      "        is_predicting, input_ids, input_mask, segment_ids, label_ids, num_labels)\r\n",
      "\r\n",
      "      train_op = optimization.create_optimizer(\r\n",
      "          loss, learning_rate, num_train_steps, num_warmup_steps, use_tpu=False)\r\n",
      "\r\n",
      "      # Calculate evaluation metrics. \r\n",
      "      def metric_fn(label_ids, predicted_labels):\r\n",
      "        accuracy = tf.metrics.accuracy(label_ids, predicted_labels)\r\n",
      "        f1_score = tf.contrib.metrics.f1_score(\r\n",
      "            label_ids,\r\n",
      "            predicted_labels)\r\n",
      "        auc = tf.metrics.auc(\r\n",
      "            label_ids,\r\n",
      "            predicted_labels)\r\n",
      "        recall = tf.metrics.recall(\r\n",
      "            label_ids,\r\n",
      "            predicted_labels)\r\n",
      "        precision = tf.metrics.precision(\r\n",
      "            label_ids,\r\n",
      "            predicted_labels) \r\n",
      "        true_pos = tf.metrics.true_positives(\r\n",
      "            label_ids,\r\n",
      "            predicted_labels)\r\n",
      "        true_neg = tf.metrics.true_negatives(\r\n",
      "            label_ids,\r\n",
      "            predicted_labels)   \r\n",
      "        false_pos = tf.metrics.false_positives(\r\n",
      "            label_ids,\r\n",
      "            predicted_labels)  \r\n",
      "        false_neg = tf.metrics.false_negatives(\r\n",
      "            label_ids,\r\n",
      "            predicted_labels)\r\n",
      "        return {\r\n",
      "            \"eval_accuracy\": accuracy,\r\n",
      "            \"f1_score\": f1_score,\r\n",
      "            \"auc\": auc,\r\n",
      "            \"precision\": precision,\r\n",
      "            \"recall\": recall,\r\n",
      "            \"true_positives\": true_pos,\r\n",
      "            \"true_negatives\": true_neg,\r\n",
      "            \"false_positives\": false_pos,\r\n",
      "            \"false_negatives\": false_neg\r\n",
      "        }\r\n",
      "\r\n",
      "      eval_metrics = metric_fn(label_ids, predicted_labels)\r\n",
      "\r\n",
      "      if mode == tf.estimator.ModeKeys.TRAIN:\r\n",
      "        return tf.estimator.EstimatorSpec(mode=mode,\r\n",
      "          loss=loss,\r\n",
      "          train_op=train_op)\r\n",
      "      else:\r\n",
      "          return tf.estimator.EstimatorSpec(mode=mode,\r\n",
      "            loss=loss,\r\n",
      "            eval_metric_ops=eval_metrics)\r\n",
      "    else:\r\n",
      "      (predicted_labels, log_probs) = create_model(\r\n",
      "        is_predicting, input_ids, input_mask, segment_ids, label_ids, num_labels)\r\n",
      "\r\n",
      "      predictions = {\r\n",
      "          'probabilities': log_probs,\r\n",
      "          'labels': predicted_labels\r\n",
      "      }\r\n",
      "      return tf.estimator.EstimatorSpec(mode, predictions=predictions)\r\n",
      "\r\n",
      "  # Return the actual model function in the closure\r\n",
      "  return model_fn\r\n",
      "\r\n",
      "\r\n",
      "def create_tokenizer_from_hub_module():\r\n",
      "    \"\"\"Get the vocab file and casing info from the Hub module.\"\"\"\r\n",
      "    with tf.Graph().as_default():\r\n",
      "        bert_module = hub.Module(BERT_MODEL_HUB)\r\n",
      "        tokenization_info = bert_module(signature=\"tokenization_info\", as_dict=True)\r\n",
      "        with tf.Session() as sess:\r\n",
      "            vocab_file, do_lower_case = sess.run([tokenization_info[\"vocab_file\"],\r\n",
      "                                                tokenization_info[\"do_lower_case\"]])\r\n",
      "      \r\n",
      "        return tokenization.FullTokenizer(vocab_file=vocab_file,\r\n",
      "                                               do_lower_case=do_lower_case)\r\n",
      "    \r\n",
      "    \r\n",
      "def predict(in_sentences):\r\n",
      "    labels = [1, 2, 3, 4, 5]\r\n",
      "\r\n",
      "    tokenizer = create_tokenizer_from_hub_module()\r\n",
      "    \r\n",
      "    input_examples = [amazon_run_classifier.InputExample(guid=\"\", text_a = x, text_b = None, label = -1) for x in in_sentences] # here, \"\" is just a dummy label\r\n",
      "\r\n",
      "    input_features = amazon_run_classifier.convert_examples_to_features(input_examples, LABEL_VALUES, MAX_SEQ_LENGTH, tokenizer)\r\n",
      "\r\n",
      "    predict_input_fn = amazon_run_classifier.input_fn_builder(features=input_features, seq_length=MAX_SEQ_LENGTH, is_training=False, drop_remainder=False)\r\n",
      "\r\n",
      "    predictions = estimator.predict(predict_input_fn)\r\n",
      "\r\n",
      "    return [(sentence, prediction['probabilities'], labels[prediction['labels']]) for sentence, prediction in zip(in_sentences, predictions)]\r\n",
      "\r\n",
      "\r\n",
      "def serving_input_fn():\r\n",
      "    label_ids = tf.placeholder(tf.int32, [None], name='label_ids')\r\n",
      "    input_ids = tf.placeholder(tf.int32, [None, MAX_SEQ_LENGTH], name='input_ids')\r\n",
      "    input_mask = tf.placeholder(tf.int32, [None, MAX_SEQ_LENGTH], name='input_mask')\r\n",
      "    segment_ids = tf.placeholder(tf.int32, [None, MAX_SEQ_LENGTH], name='segment_ids')\r\n",
      "    input_fn = tf.estimator.export.build_raw_serving_input_receiver_fn({\r\n",
      "        'label_ids': label_ids,\r\n",
      "        'input_ids': input_ids,\r\n",
      "        'input_mask': input_mask,\r\n",
      "        'segment_ids': segment_ids,\r\n",
      "    })()\r\n",
      "    return input_fn\r\n",
      "\r\n",
      "\r\n",
      "if __name__ == '__main__':\r\n",
      "    parser = argparse.ArgumentParser()\r\n",
      "\r\n",
      "#    parser.add_argument('--model-type', type=str, default='bert')\r\n",
      "#    parser.add_argument('--model-name', type=str, default='bert-base-cased')\r\n",
      "    parser.add_argument('--train-data', type=str, default=os.environ['SM_CHANNEL_TRAIN'])\r\n",
      "    parser.add_argument('--validation-data', type=str, default=os.environ['SM_CHANNEL_VALIDATION'])\r\n",
      "    parser.add_argument('--model-dir', type=str, default=os.environ['SM_MODEL_DIR'])\r\n",
      "    parser.add_argument('--hosts', type=list, default=json.loads(os.environ['SM_HOSTS']))\r\n",
      "    parser.add_argument('--current-host', type=str, default=os.environ['SM_CURRENT_HOST'])\r\n",
      "    parser.add_argument('--num-gpus', type=int, default=os.environ['SM_NUM_GPUS'])\r\n",
      "\r\n",
      "    args, _ = parser.parse_known_args()   \r\n",
      "#    model_type = args.model_type\r\n",
      "#    model_name = args.model_name\r\n",
      "    train_data = args.train_data\r\n",
      "    validation_data = args.validation_data\r\n",
      "    model_dir = args.model_dir\r\n",
      "    hosts = args.hosts\r\n",
      "    current_host = args.current_host\r\n",
      "    num_gpus = args.num_gpus\r\n",
      "\r\n",
      "    # Compute train and warmup steps from batch size\r\n",
      "    # These hyperparameters are copied from this colab notebook (https://colab.sandbox.google.com/github/tensorflow/tpu/blob/master/tools/colab/bert_finetuning_with_cloud_tpus.ipynb)\r\n",
      "    BATCH_SIZE = 32\r\n",
      "    LEARNING_RATE = 2e-5\r\n",
      "    NUM_TRAIN_EPOCHS = 3.0\r\n",
      "    # Warmup is a period of time where hte learning rate \r\n",
      "    # is small and gradually increases--usually helps training.\r\n",
      "    WARMUP_PROPORTION = 0.1\r\n",
      "    # Model configs\r\n",
      "    SAVE_CHECKPOINTS_STEPS = 500\r\n",
      "    SAVE_SUMMARY_STEPS = 100\r\n",
      "\r\n",
      "    USE_BUCKET = False \r\n",
      "\r\n",
      "    pipe_mode_str = os.environ.get('SM_INPUT_DATA_CONFIG', '')\r\n",
      "    print('pipe_mode_str {}'.format(pipe_mode_str))\r\n",
      "\r\n",
      "    pipe_mode = (pipe_mode_str.find('Pipe') >= 0)\r\n",
      "    print('pipe_mode {}'.format(pipe_mode))\r\n",
      "\r\n",
      "    # Compute # train and warmup steps from batch size\r\n",
      "    \r\n",
      "    # We need the # of training rows.  For now, just hard-coding\r\n",
      "    #num_train_steps = int(len(train_features) / BATCH_SIZE * NUM_TRAIN_EPOCHS)\r\n",
      "    num_train_steps = 100\r\n",
      "    num_warmup_steps = int(num_train_steps * WARMUP_PROPORTION)\r\n",
      "\r\n",
      "    # Specify output directory and number of checkpoint steps to save\r\n",
      "    run_config = tf.estimator.RunConfig(\r\n",
      "        model_dir=model_dir,\r\n",
      "        save_summary_steps=SAVE_SUMMARY_STEPS,\r\n",
      "        save_checkpoints_steps=SAVE_CHECKPOINTS_STEPS)\r\n",
      "\r\n",
      "    model_fn = model_fn_builder(\r\n",
      "      num_labels=len(LABEL_VALUES),\r\n",
      "      learning_rate=LEARNING_RATE,\r\n",
      "      num_train_steps=num_train_steps,\r\n",
      "      num_warmup_steps=num_warmup_steps)\r\n",
      "\r\n",
      "    estimator = tf.estimator.Estimator(\r\n",
      "      model_fn=model_fn,\r\n",
      "      config=run_config,\r\n",
      "      params={\"batch_size\": BATCH_SIZE})\r\n",
      "\r\n",
      "    # Next we create an input builder function that takes our training feature set (`train_features`) and produces a generator. This is a pretty standard design pattern for working with Tensorflow [Estimators](https://www.tensorflow.org/guide/estimators).\r\n",
      "\r\n",
      "    train_data_filenames = glob.glob('{}/*.tfrecord'.format(train_data))\r\n",
      "    print('train_data_filenames {}'.format(train_data_filenames))\r\n",
      "    \r\n",
      "    # Create an input function for training. drop_remainder = True for using TPUs.\r\n",
      "    train_input_fn = amazon_run_classifier.file_based_input_fn_builder(\r\n",
      "        input_files=train_data_filenames,\r\n",
      "        seq_length=MAX_SEQ_LENGTH,\r\n",
      "        is_training=True,\r\n",
      "        drop_remainder=False,\r\n",
      "        pipe_mode=pipe_mode)\r\n",
      "    \r\n",
      "    print('Beginning Training!')\r\n",
      "    current_time = datetime.now()\r\n",
      "    estimator.train(input_fn=train_input_fn, max_steps=num_train_steps)\r\n",
      "    print('Training took time ', datetime.now() - current_time)\r\n",
      "    print('Ending Training!')\r\n",
      "        \r\n",
      "    print('Starting Exporting!')\r\n",
      "    estimator.export_savedmodel('{}/tf-bert-model/'.format(model_dir), serving_input_fn)\r\n",
      "    print('Listing contents of {}'.format(model_dir))\r\n",
      "    model_dir = os.listdir(model_dir)\r\n",
      "    for file in model_dir:\r\n",
      "        print(file)\r\n",
      "    print('Ending Exporting!')\r\n",
      "          \r\n",
      "#   TODO:  Figure out why this gets stuck  \r\n",
      "#            tensorflow.python.framework.errors_impl.InvalidArgumentError: assertion failed: [predictions must be in [0, 1]] [Condition x <= y did not hold element-wise:] [x (f1/remove_squeezable_dimensions/cond/Merge:0) = ] [4 4 4...] [y (f1/Cast_1:0) = ] [1]  \r\n",
      "#    print('Begin Validating!')\r\n",
      "\r\n",
      "#    validation_data_filenames = glob.glob('{}/*.tfrecord'.format(validation_data))\r\n",
      "#    print('validation_data_filenames {}'.format(validation_data_filenames))\r\n",
      "#    validation_input_fn = amazon_run_classifier.file_based_input_fn_builder(\r\n",
      "#        input_files=validation_data_filenames,\r\n",
      "#        seq_length=MAX_SEQ_LENGTH,\r\n",
      "#        is_training=False,\r\n",
      "#        drop_remainder=False,\r\n",
      "#        pipe_mode=pipe_mode)\r\n",
      "\r\n",
      "#    # Now let's use our test data to see how well our model did:\r\n",
      "#    estimator.evaluate(input_fn=validation_input_fn, steps=None)\r\n",
      "\r\n",
      "#    # TODO:  Print out evaluation metrics\r\n",
      "#    print('End Validating!')\r\n",
      "    \r\n",
      "    # Now let's write code to make predictions on new sentences:\r\n",
      "    pred_sentences = [\r\n",
      "      \"This is awful.\",\r\n",
      "      \"This is just OK.\",\r\n",
      "      \"This is surprisingly creative.\",\r\n",
      "      \"This is absolutely fantastic!\",\r\n",
      "    ]\r\n",
      "\r\n",
      "    print('Begin Predicting!')\r\n",
      "    predictions = predict(pred_sentences)\r\n",
      "    print(predictions)\r\n",
      "    print('End Predicting!')\r\n",
      "\r\n",
      "    print('Complete')\r\n"
     ]
    }
   ],
   "source": [
    "!cat src_bert_tf/tf_bert_reviews.py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sagemaker.tensorflow import TensorFlow\n",
    "\n",
    "model_output_path = 's3://{}/models/tf-bert'.format(bucket)\n",
    "\n",
    "bert_estimator = TensorFlow(entry_point='tf_bert_reviews.py',\n",
    "                            source_dir='src_bert_tf',\n",
    "                            role=role,\n",
    "                            train_instance_count=1,\n",
    "                            train_instance_type='ml.p3.8xlarge',\n",
    "                            py_version='py3',\n",
    "                            framework_version='1.15.2',\n",
    "                            output_path=model_output_path,\n",
    "#                            hyperparameters={'model_type':'bert',\n",
    "#                                             'model_name': 'bert-base-cased'},\n",
    "                            distributions={'parameter_server': {'enabled': True}},\n",
    "                            enable_cloudwatch_metrics=True,\n",
    "                            input_mode='Pipe')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Train the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "metadata": {},
   "outputs": [],
   "source": [
    "bert_estimator.fit(inputs={'train': s3_input_train_data, \n",
    "                           'validation': s3_input_validation_data,                           \n",
    "                          },\n",
    "#                   compression='Gzip',\n",
    "                   wait=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 64,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "training_job_name:  tensorflow-training-2020-04-01-22-06-15-210\n"
     ]
    }
   ],
   "source": [
    "training_job_name = bert_estimator.latest_training_job.name\n",
    "print('training_job_name:  {}'.format(training_job_name))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 65,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# from sagemaker.tensorflow import TensorFlow\n",
    "\n",
    "# bert_estimator = TensorFlow.attach(training_job_name=training_job_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<b>Review <a href=\"https://console.aws.amazon.com/sagemaker/home?region=us-east-1#/jobs/tensorflow-training-2020-04-01-22-06-15-210\">Training Job</a> After About 5 Minutes</b>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from IPython.core.display import display, HTML\n",
    "\n",
    "display(HTML('<b>Review <a href=\"https://console.aws.amazon.com/sagemaker/home?region={}#/jobs/{}\">Training Job</a> After About 5 Minutes</b>'.format(region, training_job_name)))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<b>Review <a href=\"https://console.aws.amazon.com/cloudwatch/home?region=us-east-1#logStream:group=/aws/sagemaker/TrainingJobs;prefix=tensorflow-training-2020-04-01-22-06-15-210;streamFilter=typeLogStreamPrefix\">CloudWatch Logs</a> After About 5 Minutes</b>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from IPython.core.display import display, HTML\n",
    "\n",
    "display(HTML('<b>Review <a href=\"https://console.aws.amazon.com/cloudwatch/home?region={}#logStream:group=/aws/sagemaker/TrainingJobs;prefix={};streamFilter=typeLogStreamPrefix\">CloudWatch Logs</a> After About 5 Minutes</b>'.format(region, training_job_name)))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 68,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<b>Review <a href=\"https://s3.console.aws.amazon.com/s3/buckets/sagemaker-us-east-1-835319576252/models/tf-bert/tensorflow-training-2020-04-01-22-06-15-210/?region=us-east-1&tab=overview\">S3 Output Data</a> After The Training Job Has Completed</b>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from IPython.core.display import display, HTML\n",
    "\n",
    "# This is different than the job name because we are not using ProcessingOutput's in this Spark ML case.\n",
    "training_job_s3_output_prefix = 'models/tf-bert/{}'.format(training_job_name)\n",
    "\n",
    "display(HTML('<b>Review <a href=\"https://s3.console.aws.amazon.com/s3/buckets/{}/{}/?region={}&tab=overview\">S3 Output Data</a> After The Training Job Has Completed</b>'.format(bucket, training_job_s3_output_prefix, region)))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Download and Load the Trained Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "fatal error: An error occurred (404) when calling the HeadObject operation: Key \"models/tf-bert/tensorflow-training-2020-03-30-04-09-24-852/output/model.tar.gz\" does not exist\r\n"
     ]
    }
   ],
   "source": [
    "# download the model artifact from AWS S3\n",
    "!aws s3 cp $model_output_path/$training_job_name/output/model.tar.gz ./models/bert-tf/\n",
    "\n",
    "#!aws s3 cp s3://sagemaker-us-east-1-835319576252/models/tf-bert/script-mode/training-runs/tensorflow-training-2020-03-24-04-41-39-405/output/model.tar.gz ./models/bert-tf/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tarfile\n",
    "import pickle as pkl\n",
    "\n",
    "tar = tarfile.open('./models/bert-tf/model.tar.gz')\n",
    "tar.extractall(path='./models/bert-tf-model')\n",
    "tar.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "total 3958940\r\n",
      "drwxrwxr-x 5 ec2-user ec2-user       4096 Mar 30 04:09 .\r\n",
      "drwxrwxr-x 9 ec2-user ec2-user       4096 Mar 29 19:37 ..\r\n",
      "drwxr-xr-x 2 ec2-user ec2-user       4096 Mar 29 19:26 assets\r\n",
      "-rw-r--r-- 1 ec2-user ec2-user 1314587192 Mar 29 19:26 bert_reviews.h5\r\n",
      "-rw-r--r-- 1 ec2-user ec2-user        128 Mar 27 04:38 checkpoint\r\n",
      "-rw-r--r-- 1 ec2-user ec2-user   21811263 Mar 24 04:54 events.out.tfevents.1585025096.ip-10-0-168-104.ec2.internal\r\n",
      "-rw-r--r-- 1 ec2-user ec2-user   21810919 Mar 25 07:34 events.out.tfevents.1585121188.ip-10-2-248-15.ec2.internal\r\n",
      "-rw-r--r-- 1 ec2-user ec2-user   21810875 Mar 27 04:38 events.out.tfevents.1585283439.ip-10-2-66-45.ec2.internal\r\n",
      "-rw-r--r-- 1 ec2-user ec2-user   12872188 Mar 27 04:30 graph.pbtxt\r\n",
      "-rw-r--r-- 1 ec2-user ec2-user 1316323628 Mar 27 04:30 model.ckpt-0.data-00000-of-00001\r\n",
      "-rw-r--r-- 1 ec2-user ec2-user      23480 Mar 27 04:30 model.ckpt-0.index\r\n",
      "-rw-r--r-- 1 ec2-user ec2-user    5975493 Mar 27 04:30 model.ckpt-0.meta\r\n",
      "-rw-r--r-- 1 ec2-user ec2-user 1316323628 Mar 27 04:38 model.ckpt-100.data-00000-of-00001\r\n",
      "-rw-r--r-- 1 ec2-user ec2-user      23480 Mar 27 04:38 model.ckpt-100.index\r\n",
      "-rw-r--r-- 1 ec2-user ec2-user    5975493 Mar 27 04:38 model.ckpt-100.meta\r\n",
      "-rw-r--r-- 1 ec2-user ec2-user   16356565 Mar 29 19:26 saved_model.pb\r\n",
      "drwxr-xr-x 3 ec2-user ec2-user       4096 Mar 27 04:38 tf-bert-model-oh-yeah\r\n",
      "drwxr-xr-x 2 ec2-user ec2-user       4096 Mar 29 19:26 variables\r\n"
     ]
    }
   ],
   "source": [
    "!ls -al ./models/bert-tf-model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "total 2076\r\n",
      "drwxr-xr-x 4 ec2-user ec2-user    4096 Mar 27 04:38 .\r\n",
      "drwxr-xr-x 3 ec2-user ec2-user    4096 Mar 27 04:38 ..\r\n",
      "drwxr-xr-x 2 ec2-user ec2-user    4096 Mar 27 04:38 assets\r\n",
      "-rw-r--r-- 1 ec2-user ec2-user 2107865 Mar 27 04:38 saved_model.pb\r\n",
      "drwxr-xr-x 2 ec2-user ec2-user    4096 Mar 27 04:38 variables\r\n"
     ]
    }
   ],
   "source": [
    "!ls -al ./models/bert-tf-model/tf-bert/1585283912"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[33mYou are using pip version 10.0.1, however version 20.0.2 is available.\n",
      "You should consider upgrading via the 'pip install --upgrade pip' command.\u001b[0m\n",
      "\u001b[31m  Could not find a version that satisfies the requirement tensorflow==1.15.2 (from versions: 0.12.1, 1.0.0, 1.0.1, 1.1.0rc0, 1.1.0rc1, 1.1.0rc2, 1.1.0, 1.2.0rc0, 1.2.0rc1, 1.2.0rc2, 1.2.0, 1.2.1, 1.3.0rc0, 1.3.0rc1, 1.3.0rc2, 1.3.0, 1.4.0rc0, 1.4.0rc1, 1.4.0, 1.4.1, 1.5.0rc0, 1.5.0rc1, 1.5.0, 1.5.1, 1.6.0rc0, 1.6.0rc1, 1.6.0, 1.7.0rc0, 1.7.0rc1, 1.7.0, 1.7.1, 1.8.0rc0, 1.8.0rc1, 1.8.0, 1.9.0rc0, 1.9.0rc1, 1.9.0rc2, 1.9.0, 1.10.0rc0, 1.10.0rc1, 1.10.0, 1.10.1, 1.11.0rc0, 1.11.0rc1, 1.11.0rc2, 1.11.0, 1.12.0rc0, 1.12.0rc1, 1.12.0rc2, 1.12.0, 1.12.2, 1.12.3, 1.13.0rc0, 1.13.0rc1, 1.13.0rc2, 1.13.1, 1.13.2, 1.14.0rc0, 1.14.0rc1, 1.14.0, 2.0.0a0, 2.0.0b0, 2.0.0b1)\u001b[0m\n",
      "\u001b[31mNo matching distribution found for tensorflow==1.15.2\u001b[0m\n",
      "\u001b[33mYou are using pip version 10.0.1, however version 20.0.2 is available.\n",
      "You should consider upgrading via the 'pip install --upgrade pip' command.\u001b[0m\n",
      "\u001b[33mYou are using pip version 10.0.1, however version 20.0.2 is available.\n",
      "You should consider upgrading via the 'pip install --upgrade pip' command.\u001b[0m\n",
      "\u001b[33mYou are using pip version 10.0.1, however version 20.0.2 is available.\n",
      "You should consider upgrading via the 'pip install --upgrade pip' command.\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "# Must upgrade wrapt before installing TF\n",
    "!pip install -q wrapt --upgrade --ignore-installed\n",
    "!pip install -q tensorflow==1.15.2\n",
    "!pip install -q tensorflow-hub==0.7.0\n",
    "!pip install -q bert-tensorflow==1.0.1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Load the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'tf' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-19-3cfacfe4f1c7>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m()\u001b[0m\n\u001b[0;32m----> 1\u001b[0;31m saved_model = tf.saved_model.load_v2(\n\u001b[0m\u001b[1;32m      2\u001b[0m     \u001b[0;34m'./models/bert-tf-model/tf-bert-model-oh-yeah/1585283912'\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      3\u001b[0m     \u001b[0mtags\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mNone\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      4\u001b[0m )\n",
      "\u001b[0;31mNameError\u001b[0m: name 'tf' is not defined"
     ]
    }
   ],
   "source": [
    "saved_model = tf.saved_model.load_v2(\n",
    "    './models/bert-tf-model/tf-bert-model-oh-yeah/1585283912',\n",
    "    tags=None\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "inference = saved_model.signatures[\"serving_default\"]\n",
    "print(inference.inputs)\n",
    "print(inference.structured_outputs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Predict \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from bert import tokenization\n",
    "import tensorflow_hub as hub\n",
    "\n",
    "BERT_MODEL_HUB = \"https://tfhub.dev/google/bert_uncased_L-12_H-768_A-12/1\"\n",
    "\n",
    "def create_tokenizer_from_hub_module():\n",
    "    \"\"\"Get the vocab file and casing info from the Hub module.\"\"\"\n",
    "    with tf.Graph().as_default():\n",
    "        bert_module = hub.Module(BERT_MODEL_HUB)\n",
    "        tokenization_info = bert_module(signature=\"tokenization_info\", as_dict=True)\n",
    "        with tf.Session() as sess:\n",
    "            vocab_file, do_lower_case = sess.run([tokenization_info[\"vocab_file\"],\n",
    "                                                tokenization_info[\"do_lower_case\"]])\n",
    "      \n",
    "        return tokenization.FullTokenizer(vocab_file=vocab_file,\n",
    "                                               do_lower_case=do_lower_case)\n",
    "   "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_predict_features(features, seq_length):\n",
    "  all_input_ids = []\n",
    "  all_input_mask = []\n",
    "  all_segment_ids = []\n",
    "  all_label_ids = []\n",
    "\n",
    "  for feature in features:\n",
    "    all_input_ids.append(feature.input_ids)\n",
    "    all_input_mask.append(feature.input_mask)\n",
    "    all_segment_ids.append(feature.segment_ids)\n",
    "    all_label_ids.append(feature.label_id)\n",
    "\n",
    "    batch_size = 32\n",
    "\n",
    "  num_examples = len(features)\n",
    "\n",
    "  # This is for demo purposes and does NOT scale to large data sets. We do\n",
    "  # not use Dataset.from_generator() because that uses tf.py_func which is\n",
    "  # not TPU compatible. The right way to load data is with TFRecordReader.\n",
    "#   d = tf.data.Dataset.from_tensor_slices({\n",
    "#     \"input_ids\":\n",
    "#         tf.constant(\n",
    "#             all_input_ids, shape=[num_examples, seq_length],\n",
    "#             dtype=tf.int32),\n",
    "#     \"input_mask\":\n",
    "#         tf.constant(\n",
    "#             all_input_mask,\n",
    "#             shape=[num_examples, seq_length],\n",
    "#             dtype=tf.int32),\n",
    "#     \"segment_ids\":\n",
    "#         tf.constant(\n",
    "#             all_segment_ids,\n",
    "#             shape=[num_examples, seq_length],\n",
    "#             dtype=tf.int32),\n",
    "#     \"label_ids\":\n",
    "#         tf.constant(all_label_ids, shape=[num_examples], dtype=tf.int32),\n",
    "#   })\n",
    "\n",
    "#   d = d.batch(batch_size=batch_size, drop_remainder=False)\n",
    "\n",
    "#  return d\n",
    "\n",
    "  input_ids = tf.constant(\n",
    "             all_input_ids, shape=[num_examples, seq_length],\n",
    "             dtype=tf.int32)\n",
    "\n",
    "  input_mask = tf.constant(\n",
    "             all_input_mask,\n",
    "             shape=[num_examples, seq_length],\n",
    "             dtype=tf.int32)\n",
    "\n",
    "  segment_ids = tf.constant(\n",
    "             all_segment_ids,\n",
    "             shape=[num_examples, seq_length],\n",
    "             dtype=tf.int32)\n",
    "\n",
    "  label_ids = tf.constant(all_label_ids, shape=[num_examples], dtype=tf.int32),\n",
    "\n",
    "  return input_ids, input_mask, segment_ids, label_ids"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from src_bert_tf import amazon_run_classifier\n",
    "\n",
    "MAX_SEQ_LENGTH = 128\n",
    "LABEL_VALUES = [1, 2, 3, 4, 5]\n",
    "\n",
    "def predict(in_sentences):\n",
    "    tokenizer = create_tokenizer_from_hub_module()\n",
    "    print('**** TOKENIZER {}****'.format(tokenizer))\n",
    "    \n",
    "    input_examples = [amazon_run_classifier.InputExample(guid=\"\", text_a = x, text_b = None, label = -1) for x in in_sentences]\n",
    "    input_features = amazon_run_classifier.convert_examples_to_features(input_examples, LABEL_VALUES, MAX_SEQ_LENGTH, tokenizer)\n",
    "\n",
    "#    predict_input_fn = amazon_run_classifier.input_fn_builder(features=input_features, seq_length=MAX_SEQ_LENGTH, is_training=False, drop_remainder=False)\n",
    "#    predictions = estimator.predict(predict_input_fn)\n",
    "    input_ids, input_mask, segment_ids, label_ids = get_predict_features(input_features, MAX_SEQ_LENGTH)\n",
    "    print(type(input_ids))\n",
    "\n",
    "    inference(input_ids=input_ids, input_mask=input_mask, segment_ids=segment_ids, label_ids=label_ids)\n",
    "\n",
    "    return [(sentence, prediction['probabilities'], LABEL_VALUES[prediction['labels']]) for sentence, prediction in zip(in_sentences, predictions)]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "pred_sentences = [\n",
    "  \"That movie was absolutely awful\",\n",
    "  \"The acting was a bit lacking\",\n",
    "  \"The film was creative and surprising\",\n",
    "  \"Absolutely fantastic!\"\n",
    "]\n",
    "\n",
    "np_list = np.asarray(pred_sentences)\n",
    "tensor_list = tf.convert_to_tensor(np_list)\n",
    "predictions = predict(tensor_list)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "conda_python3",
   "language": "python",
   "name": "conda_python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
